package com.lkx.kafka.monitor.common;

import com.lkx.kafka.monitor.dto.ThroughputMetrics;
import com.lkx.kafka.monitor.service.DataStorageService;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.admin.*;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

@Component
@Slf4j
public class KafkaMetricsCollector {

    @Autowired
    private DataStorageService dataStorageService;


    private final Map<String, Map<String, Long>> lastOffsets = new ConcurrentHashMap<>();
    private final Map<String, Map<String, Long>> lastTimestamps = new ConcurrentHashMap<>();


    @Autowired
    private AdminClient adminClient;
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(2);

    @PostConstruct
    public void init() {
        // 启动定时收集任务
        scheduler.scheduleAtFixedRate(this::collectMetrics, 0, 30, TimeUnit.SECONDS);
    }



    /**
     * 收集Kafka生产和消费指标
     */
    private void collectMetrics() {
        try {

            // 获取所有主题
            Set<String> topics = adminClient.listTopics().names().get(10, TimeUnit.SECONDS);

            long currentTime = System.currentTimeMillis();
            long totalProduceRate = 0;
            long totalConsumeRate = 0;

            // 为每个主题收集指标
            for (String topic : topics) {
                if (topic.startsWith("__")) continue; // 跳过内部主题

                ThroughputMetrics metrics = collectTopicMetrics(topic, currentTime);
                if (metrics != null) {
                    dataStorageService.saveData(metrics);
                    totalProduceRate += metrics.getProduceRate() != null ? metrics.getProduceRate() : 0;
                    totalConsumeRate += metrics.getConsumeRate() != null ? metrics.getConsumeRate() : 0;
                }
            }

            // 保存聚合数据（所有主题）
            ThroughputMetrics aggregateMetrics = new ThroughputMetrics();
            aggregateMetrics.setTopicName("ALL");
            aggregateMetrics.setTimestamp(currentTime);
            aggregateMetrics.setProduceRate(totalProduceRate);
            aggregateMetrics.setConsumeRate(totalConsumeRate);
            aggregateMetrics.setTimeDimension("realtime");
            dataStorageService.saveData(aggregateMetrics);

            log.info("成功收集指标 - 生产速率: {}, 消费速率: {}", totalProduceRate, totalConsumeRate);
        } catch (Exception e) {
            log.error("收集Kafka指标失败: {}", e.getMessage());
        }
    }

    /**
 * 收集单个主题的指标
 */
private ThroughputMetrics collectTopicMetrics(String topicName, long currentTime) {
    if (topicName == null || topicName.isEmpty()) {
        log.warn("无效的主题名称");
        return null;
    }

    if (currentTime < 0) {
        log.warn("非法的时间戳: {}", currentTime);
        return null;
    }

    try {
        ThroughputMetrics metrics = new ThroughputMetrics();
        metrics.setTopicName(topicName);
        metrics.setTimestamp(currentTime);
        metrics.setTimeDimension("realtime");

        // 获取主题描述信息
        DescribeTopicsResult describeTopicsResult = adminClient.describeTopics(Collections.singletonList(topicName));
        TopicDescription topicDescription = describeTopicsResult.values().get(topicName).get(10, TimeUnit.SECONDS);

        // 批量获取所有分区的最新位移
        List<TopicPartition> partitions = topicDescription.partitions().stream()
                .map(p -> new TopicPartition(topicName, p.partition()))
                .collect(Collectors.toList());

        Map<TopicPartition, OffsetSpec> offsetSpecMap = partitions.stream()
                .collect(Collectors.toMap(tp -> tp, tp -> OffsetSpec.latest()));

        ListOffsetsResult listOffsetsResult = adminClient.listOffsets(offsetSpecMap);
        long endOffsetSum = 0;

        for (TopicPartition partition : partitions) {
            ListOffsetsResult.ListOffsetsResultInfo info = listOffsetsResult.partitionResult(partition).get();
            endOffsetSum += info.offset();
        }

        metrics.setEndOffset(endOffsetSum);

        // 计算当前位移（消费速率基于消费者组位移变化）
        long currentOffsetSum = 0;
        try {
            Collection<ConsumerGroupListing> groups = adminClient.listConsumerGroups()
                    .valid().get(10, TimeUnit.SECONDS);

            for (ConsumerGroupListing group : groups) {
                try {
                    Map<TopicPartition, OffsetAndMetadata> offsets = adminClient
                            .listConsumerGroupOffsets(group.groupId())
                            .partitionsToOffsetAndMetadata().get(10, TimeUnit.SECONDS);

                    for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : offsets.entrySet()) {
                        if (entry.getKey().topic().equals(topicName)) {
                            currentOffsetSum += entry.getValue().offset();
                        }
                    }
                } catch (Exception innerEx) {
                    log.warn("获取消费者组 {} 的位移信息失败", group.groupId(), innerEx);
                }
            }
        } catch (Exception e) {
            log.warn("获取消费者组列表失败", e);
        }

        metrics.setCurrentOffset(currentOffsetSum);

        // 计算生产速率和消费速率（简化实现，实际应该基于历史数据计算变化率）
        metrics.setProduceRate(calculateRate(topicName, "produce", endOffsetSum, currentTime));
        metrics.setConsumeRate(calculateRate(topicName, "consume", currentOffsetSum, currentTime));

        return metrics;
    } catch (Exception e) {
        log.error("收集主题 {} 指标失败", topicName, e);
        return null;
    }
}


    // 简化的速率计算方法（实际应基于历史数据）
    private Long calculateRate(String topicName, String type, long currentOffset, long currentTime) {
        try {
            // 获取上次记录的偏移量和时间戳
            Map<String, Long> lastOffsetMap = lastOffsets.computeIfAbsent(type, k -> new ConcurrentHashMap<>());
            Map<String, Long> lastTimestampMap = lastTimestamps.computeIfAbsent(type, k -> new ConcurrentHashMap<>());

            Long lastOffset = lastOffsetMap.get(topicName);
            Long lastTimestamp = lastTimestampMap.get(topicName);

            // 更新当前值
            lastOffsetMap.put(topicName, currentOffset);
            lastTimestampMap.put(topicName, currentTime);

            // 如果有历史数据，则计算速率
            if (lastOffset != null && lastTimestamp != null && currentTime > lastTimestamp) {
                long offsetDiff = currentOffset - lastOffset;
                long timeDiffSeconds = (currentTime - lastTimestamp) / 1000; // 转换为秒

                if (timeDiffSeconds > 0) {
                    return offsetDiff / timeDiffSeconds; // 每秒的消息数量
                }
            }

            // 没有足够数据计算速率时返回0
            return 0L;
        } catch (Exception e) {
            log.warn("计算主题 {} 的 {} 速率失败", topicName, type, e);
            return 0L;
        }
    }

    @PreDestroy
    public void destroy() {
        if (adminClient != null) {
            adminClient.close();
        }
        scheduler.shutdown();
    }



}
